﻿using System;
using System.Security.Cryptography;
using Microsoft.VisualStudio.TestTools.UnitTesting;
using Rhino.Mocks;
using SjclHelpers.Codec;
using SjclHelpers.Mode;

namespace SjclHelpersTests.Mode
{
    [TestClass]
    public class CcmCipherTextTests
    {
        public byte[] ValidCipherText
        {
            get { return new byte[0]; }
        }

        public byte[] ValidNonce
        {
            get { return "009e13fcac8f82e6".ToBytes(); }
        }

        public int ValidMacLength
        {
            get { return 4; }
        }

        [TestMethod]
        public void Ctor_Works()
        {
            // Fixture setup...
            const string encryptionKey = "37fceaa7744701bc58320cf440396871";
            var aesBlockCipher = new AesManaged
            {
                Key = encryptionKey.ToBytes(),
                BlockSize = 128
            };
            const string cipherText = "ec913dbff7";
            const int expectedMessageAuthenticationCodeLength = 4;  // 32 bit
            const string nonce = "009e13fcac8f82e6";

            // Exercise SUT...
            var ccmCipherText = new CcmCipherText(
                aesBlockCipher,
                cipherText.ToBytes(),
                nonce.ToBytes(),
                expectedMessageAuthenticationCodeLength
            );

            // Verify SUT...
            Assert.AreEqual(
                aesBlockCipher,
                ccmCipherText.AesBlockCipher
            );
            CollectionAssert.AreEqual(
                cipherText.ToBytes(),
                ccmCipherText.CipherText
            );
            Assert.AreEqual(
                expectedMessageAuthenticationCodeLength,
                ccmCipherText.MessageAuthenticationCodeLength
            );
            CollectionAssert.AreEqual(
                nonce.ToBytes(),
                ccmCipherText.Nonce
            );
        }

        [TestMethod]
        public void GetKeyStreamBlock_Index0_CorrectBlock()
        {
            // Fixture setup...
            const string nonce = "009e13fcac8f82e6";
            const string expectedKeyStreamBlock =
                "06009e13fcac8f82e600000000000000";
            var ccmCipherTextMock = MockRepository.GeneratePartialMock<CcmCipherText>(
                MockRepository.GenerateMock<Aes>(),
                new byte[0],
                nonce.ToBytes(),
                4
            );
            ccmCipherTextMock
                .Expect(c => c.LengthFieldSize)
                .Return(7);
            ccmCipherTextMock
                .Expect(c => c.Nonce)
                .Return(nonce.ToBytes());

            // Exercise SUT...
            var actualKeyStreamBlock = ccmCipherTextMock.GetKeyStreamBlock(0);

            // Verify SUT...
            ccmCipherTextMock.VerifyAllExpectations();
            CollectionAssert.AreEqual(
                expectedKeyStreamBlock.ToBytes(),
                actualKeyStreamBlock
            );
        }

        [TestMethod]
        public void GetEncryptedKeyStreamBlock_Index0_CorrectBlock()
        {
            // Fixture setup...
            const string encryptionKey = "37fceaa7744701bc58320cf440396871";
            var aesBlockCipher = new AesManaged
            {
                Key = encryptionKey.ToBytes(),
                BlockSize = 128,
                Mode = CipherMode.ECB
            };
            const string nonce = "009e13fcac8f82e6";
            const string keyStreamBlock =
                "06009e13fcac8f82e600000000000000";
            const string expectedEncryptedKeyStreamBlock =
                "3bf57b8f3fe0bd0ccfef0bc841ccf51c";
            var ccmCipherTextMock = MockRepository
                .GeneratePartialMock<CcmCipherText>(
                    MockRepository.GenerateMock<Aes>(),
                    new byte[0], nonce.ToBytes(), 4
                );
            ccmCipherTextMock
                .Expect(c => c.AesBlockCipher)
                .Return(aesBlockCipher);
            ccmCipherTextMock
                .Expect(c => c.GetKeyStreamBlock(0))
                .Return(keyStreamBlock.ToBytes());

            // Exercise SUT...
            var actualEncryptedKeyStreamBlock =
                ccmCipherTextMock.GetEncryptedKeyStreamBlock(0);

            // Verify SUT...
            ccmCipherTextMock.VerifyAllExpectations();
            CollectionAssert.AreEqual(
                expectedEncryptedKeyStreamBlock.ToBytes(),
                actualEncryptedKeyStreamBlock
            );
        }

        [TestMethod]
        public void MessageAuthenticationCode_Works()
        {
            // Fixture setup...
            const string cipherText = "ec913dbff7";
            const string expectedMac = "aac8c478";
            const string encryptedKeyStreamBlock =
                "3bf57b8f3fe0bd0ccfef0bc841ccf51c";
            var ccmCipherTextMock = MockRepository
                .GeneratePartialMock<CcmCipherText>(
                    MockRepository.GenerateMock<Aes>(),
                    new byte[0], ValidNonce, 4
                );
            ccmCipherTextMock
                .Expect(c => c.MessageAuthenticationCodeLength)
                .Return(4);
            ccmCipherTextMock
                .Expect(c => c.MessageLength)
                .Return(1);
            ccmCipherTextMock
                .Expect(c => c.CipherText)
                .Return(cipherText.ToBytes());
            ccmCipherTextMock
                .Expect(c => c.GetEncryptedKeyStreamBlock(0))
                .Return(encryptedKeyStreamBlock.ToBytes());

            // Exercise SUT...
            var actualMac =
                ccmCipherTextMock.MessageAuthenticationCode;

            // Verify SUT...
            ccmCipherTextMock.VerifyAllExpectations();
            CollectionAssert.AreEqual(expectedMac.ToBytes(), actualMac);
        }

        [TestMethod]
        public void MessageLength_Works()
        {
            // Fixture setup...
            var cipherText = new byte[104];
            const int messageAuthenticationCodeLength = 4;
            var ccmCipherTextMock = MockRepository
                .GeneratePartialMock<CcmCipherText>(
                    MockRepository.GenerateMock<Aes>(),
                    cipherText, ValidNonce, messageAuthenticationCodeLength
                );
            ccmCipherTextMock
                .Expect(c => c.MessageAuthenticationCodeLength)
                .Return(messageAuthenticationCodeLength);
            ccmCipherTextMock
                .Expect(c => c.CipherText)
                .Return(cipherText);

            // Exercise SUT...
            var actualMessageLength = ccmCipherTextMock.MessageLength;

            // Verify SUT...
            ccmCipherTextMock.VerifyAllExpectations();
            Assert.AreEqual(100, actualMessageLength);
        }

        [TestMethod]
        public void LengthFieldSize_MinLength_2()
        {
            // Fixture setup...
            var nonce = new byte[13];
            var ccmCipherTextMock = MockRepository
                .GeneratePartialMock<CcmCipherText>(
                    MockRepository.GenerateMock<Aes>(),
                    ValidCipherText, ValidNonce, ValidMacLength
                );
            ccmCipherTextMock
                .Expect(c => c.Nonce)
                .Return(nonce);

            // Exercise SUT...
            var actualLengthFieldSize = ccmCipherTextMock.LengthFieldSize;

            // Verify SUT...
            ccmCipherTextMock.VerifyAllExpectations();
            Assert.AreEqual(2, actualLengthFieldSize);
        }

        [TestMethod]
        [ExpectedException(typeof(InvalidOperationException))]
        public void LengthFieldSize_NonceToBig_InvalidOperationException()
        {
            // Fixture setup...
            var nonce = new byte[14];
            var ccmCipherTextMock = MockRepository
                .GeneratePartialMock<CcmCipherText>(
                    MockRepository.GenerateMock<Aes>(),
                    ValidCipherText, ValidNonce, ValidMacLength
                );
            ccmCipherTextMock
                .Expect(c => c.Nonce)
                .Return(nonce);

            // Exercise SUT...
            var dummy = ccmCipherTextMock.LengthFieldSize;
        }

        [TestMethod]
        [ExpectedException(typeof(InvalidOperationException))]
        public void LengthFieldSize_NonceToSmall_InvalidOperationException()
        {
            // Fixture setup...
            var nonce = new byte[6];
            var ccmCipherTextMock = MockRepository
                .GeneratePartialMock<CcmCipherText>(
                    MockRepository.GenerateMock<Aes>(),
                    ValidCipherText, ValidNonce, ValidMacLength
                );
            ccmCipherTextMock
                .Expect(c => c.Nonce)
                .Return(nonce);

            // Exercise SUT...
            var dummy = ccmCipherTextMock.LengthFieldSize;
        }

        [TestMethod]
        [ExpectedException(typeof(InvalidOperationException))]
        public void AuthenticationField_MacToShort_InvalidOperationException()
        {
            // Fixture setup...
            const int invalidMacLength = 3;
            var ccmCipherTextMock = MockRepository
                .GeneratePartialMock<CcmCipherText>(
                    MockRepository.GenerateMock<Aes>(),
                    ValidCipherText, ValidNonce, ValidMacLength
                );
            ccmCipherTextMock
                .Expect(c => c.MessageAuthenticationCodeLength)
                .Return(invalidMacLength);

            // Exercise SUT...
            var dummy = ccmCipherTextMock.AuthenticationField;
        }

        [TestMethod]
        [ExpectedException(typeof(InvalidOperationException))]
        public void AuthenticationField_MacToLong_InvalidOperationException()
        {
            // Fixture setup...
            const int invalidMacLength = 17;
            var ccmCipherTextMock = MockRepository
                .GeneratePartialMock<CcmCipherText>(
                    MockRepository.GenerateMock<Aes>(),
                    ValidCipherText, ValidNonce, ValidMacLength
                );
            ccmCipherTextMock
                .Expect(c => c.MessageAuthenticationCodeLength)
                .Return(invalidMacLength)
                .Repeat.Any();

            // Exercise SUT...
            var dummy = ccmCipherTextMock.AuthenticationField;
        }

        [TestMethod]
        [ExpectedException(typeof(InvalidOperationException))]
        public void AuthenticationField_OddMacLength_InvalidOperationException()
        {
            // Fixture setup...
            const int invalidMacLength = 5;
            var ccmCipherTextMock = MockRepository
                .GeneratePartialMock<CcmCipherText>(
                    MockRepository.GenerateMock<Aes>(),
                    ValidCipherText, ValidNonce, ValidMacLength
                );
            ccmCipherTextMock
                .Expect(c => c.MessageAuthenticationCodeLength)
                .Return(invalidMacLength)
                .Repeat.Any();

            // Exercise SUT...
            var dummy = ccmCipherTextMock.AuthenticationField;
        }

        [TestMethod]
        public void AuthenticationField_MaxMacLength_ValidField()
        {
            // Fixture setup...
            var expectedAuthenticationField =
                Convert.ToString(7, 2).PadLeft(3, '0');
            const int maxMacLength = 16;
            var ccmCipherTextMock = MockRepository
                .GeneratePartialMock<CcmCipherText>(
                    MockRepository.GenerateMock<Aes>(),
                    ValidCipherText, ValidNonce, ValidMacLength
                );
            ccmCipherTextMock
                .Expect(c => c.MessageAuthenticationCodeLength)
                .Return(maxMacLength)
                .Repeat.Any();

            // Exercise SUT...
            var actualAuthenticationField =
                ccmCipherTextMock.AuthenticationField;

            // Verify SUT...
            Assert.AreEqual(
                expectedAuthenticationField,
                actualAuthenticationField
            );
        }

        [TestMethod]
        public void AuthenticationField_MinMacLength_ValidField()
        {
            // Fixture setup...
            var expectedAuthenticationField =
                Convert.ToString(1, 2).PadLeft(3, '0');
            const int minMacLength = 4;
            var ccmCipherTextMock = MockRepository
                .GeneratePartialMock<CcmCipherText>(
                    MockRepository.GenerateMock<Aes>(),
                    ValidCipherText, ValidNonce, ValidMacLength
                );
            ccmCipherTextMock
                .Expect(c => c.MessageAuthenticationCodeLength)
                .Return(minMacLength)
                .Repeat.Any();

            // Exercise SUT...
            var actualAuthenticationField =
                ccmCipherTextMock.AuthenticationField;

            // Verify SUT...
            Assert.AreEqual(
                expectedAuthenticationField,
                actualAuthenticationField
            );
        }

        //
        //
        // TODO: Isolated InitializationVector tests...
        //
        //

        [TestMethod]
        public void InitializationVector_Works()
        {
            // Fixture setup...
            const int macLength = 4;  // 32 bit
            const string cipherText = "ec913dbff7";
            const string nonce = "009e13fcac8f82e6";
            const string expectedInitializationVector =
                "0e009e13fcac8f82e600000000000001";
            var expectedAuthenticationField =
                Convert.ToString(1, 2).PadLeft(3, '0');
            var ccmCipherText = new CcmCipherText(
                MockRepository.GenerateMock<Aes>(),
                cipherText.ToBytes(), nonce.ToBytes(), macLength
            );

            // Exercise SUT...
            var actualInitializationVector =
                ccmCipherText.InitializationVector;

            // Verify SUT...
            CollectionAssert.AreEqual(
                expectedInitializationVector.ToBytes(),
                actualInitializationVector
            );
        }
    }
}